# -*- coding: utf-8 -*-
"""
项目名称：cotton_test
文件名称：dataset.py
创建时间：2025-06-08 20:47:50

系统用户：wyl
作　　者：無以菱
联系邮箱：huangjing2001.guet@qq.com
功能描述：棉花叶片数据集加载模块
- 本模块定义了用于加载和处理棉花叶片图像数据集的Dataset类。
- 该类负责从磁盘读取图像文件，应用转换，并提供批量加载接口。
"""

import os
import torch
from torch.utils.data import Dataset
from PIL import Image

class CottonLeafDataset(Dataset):
    """
    棉花叶片数据集类
    
    该类继承自PyTorch的Dataset类，用于加载棉花叶片健康/患病分类数据集。
    数据集结构应为：
    root_dir/
        fresh_cotton_leaf/
            image1.jpg
            image2.jpg
            ...
        diseased_cotton_leaf/
            image1.jpg
            image2.jpg
            ...
    """
    def __init__(self, root_dir, transform=None):
        """
        初始化棉花叶片数据集
        
        参数:
            root_dir (string): 数据集根目录
            transform (callable, optional): 可选的图像转换操作
        """
        self.root_dir = root_dir
        self.transform = transform
        self.images = []
        self.labels = []
        self.label_map = {}
        
        # 直接设置标签映射，避免依赖目录遍历顺序
        # 将类别名称映射到数字标签：健康叶片->0，患病叶片->1
        self.label_map = {
            "fresh_cotton_leaf": 0,      # 健康的棉花叶片
            "diseased_cotton_leaf": 1    # 患病的棉花叶片
        }
        
        print(f"正在加载数据集，根目录: {root_dir}")
        print(f"查找的文件夹: {list(self.label_map.keys())}")
        
        # 遍历数据集目录，加载所有图像和对应的标签
        for label_name, label_idx in self.label_map.items():
            # 构建完整的类别目录路径
            label_dir = os.path.join(self.root_dir, label_name)
            
            # 检查目录是否存在
            if not os.path.exists(label_dir):
                print(f"警告: 文件夹 {label_dir} 不存在!")
                continue
                
            print(f"处理文件夹: {label_dir}")
            
            # 遍历该类别目录下的所有图像文件
            for file in os.listdir(label_dir):
                # 只处理图像文件（PNG、JPG、JPEG格式）
                if file.endswith(('.png', '.jpg', '.jpeg')):
                    # 构建完整的图像文件路径
                    file_path = os.path.join(label_dir, file)
                    # 将图像路径添加到列表
                    self.images.append(file_path)
                    # 将对应的类别标签添加到列表
                    self.labels.append(label_idx)
        
        print(f"数据集加载完成，共 {len(self.images)} 张图像")
        print(f"类别分布: {self._get_class_distribution()}")
    
    def _get_class_distribution(self):
        """
        统计数据集中各类别的样本数量
        
        返回:
            dict: 包含每个类别及其对应样本数量的字典
        """
        # 初始化计数字典
        class_counts = {}
        
        # 遍历所有标签，统计每个类别的样本数量
        for label in self.labels:
            if label in class_counts:
                class_counts[label] += 1  # 已有类别，计数加1
            else:
                class_counts[label] = 1   # 新类别，初始化计数为1
        
        # 将数字标签转换回可读的类别名称
        result = {}
        # 创建标签值到类别名称的反向映射
        label_to_name = {v: k for k, v in self.label_map.items()}
        
        # 遍历计数结果，将数字标签替换为类别名称
        for label_idx, count in class_counts.items():
            # 获取类别名称，如果找不到则显示为"未知类别"
            result[label_to_name.get(label_idx, f"未知类别{label_idx}")] = count
        
        return result
    
    def __len__(self):
        """
        返回数据集中的样本数量
        """
        return len(self.images)
    
    def __getitem__(self, idx):
        """
        获取指定索引的样本
        
        此方法在数据加载器调用时被执行，用于获取批次中的单个样本。
        
        参数:
            idx (int): 样本索引
            
        返回:
            tuple: (图像张量, 标签) - 处理后的图像张量和对应的类别标签
        """
        # 获取图像文件路径
        img_path = self.images[idx]
        
        # 使用PIL库打开图像文件并转换为RGB格式
        # 转换为RGB确保图像有3个通道，即使原图是灰度或其他格式
        image = Image.open(img_path).convert("RGB")
        
        # 获取对应的标签
        label = self.labels[idx]
        
        # 如果提供了转换函数，则应用转换
        # 转换通常包括调整大小、裁剪、归一化等操作
        if self.transform:
            image = self.transform(image)
            
        # 返回(图像, 标签)元组
        return image, label 